# -*- coding:utf-8 -*-
"""
该脚本将从Houdini中导出的abc材质导入到maya,配合材质文件将其添加到maya中
"""
import os
import sys
from PySide2.QtGui import *
from PySide2.QtCore import *
from PySide2.QtWidgets import *

try:
    import maya.cmds as cmds
    from maya import OpenMayaUI as omui
    from shiboken2 import wrapInstance
except Exception as import_err:
    print('test')


def maya_main_window():
    """
     Return the Maya main window widget as a Python object
    :return:
    """
    main_window_ptr = omui.MQtUtil.mainWindow()
    return wrapInstance(long(main_window_ptr), QWidget)


class MapTable(QTableWidget):

    def __init__(self):
        super(MapTable, self).__init__()
        self.data_list = []
        self.init_table()

    def init_table(self):
        self.setColumnCount(2)
        self.h_list = [u'物体名称', u'材质名称']
        self.setHorizontalHeaderLabels(self.h_list)
        self.horizontalHeader().setDefaultSectionSize(170)

        self.setRowCount(100)

    def set_data(self, data):
        self.clear()
        self.init_table()
        if data:
            if len(data) > 100:
                self.setRowCount(len(data))
            keys = data.keys()
            for r in range(len(keys)):
                mat_name = data[keys[r]]
                mat_item = QTableWidgetItem()
                mat_item.setText(mat_name)
                mat_item.setTextAlignment(Qt.AlignCenter)
                obj_item = QTableWidgetItem()
                obj_item.setText(keys[r])
                obj_item.setTextAlignment(Qt.AlignCenter)
                self.setItem(r, 0, obj_item)
                self.setItem(r, 1, mat_item)

    def get_map(self):
        mat_map = {}
        for r in range(self.rowCount()):
            obj_item = self.item(r, 0)
            mat_item = self.item(r, 1)
            obj_name = obj_item.text() if obj_item else ''
            mat_name = mat_item.text() if mat_item else ''
            if obj_name:
                mat_map[obj_name] = mat_name
        # print(mat_map)
        return mat_map


class SetMat(QDialog):

    def __init__(self, parent=None):
        super(SetMat, self).__init__(parent)

        self.setWindowTitle(u'赋材质')
        self.resize(410, 500)
        self.mid_font = QFont('', 15, 65)
        self.main_layout = QVBoxLayout()
        self.main_layout.setSpacing(15)
        self.main_title = QLabel(u'材质名称映射')
        self.main_title.setFont(self.mid_font)
        self.main_title.setAlignment(Qt.AlignCenter)
        self.mat_map_layout = QHBoxLayout()
        self.map_label = QLabel(u'MatMap路径:')
        self.map_edit = QLineEdit()
        self.map_choose_btn = QPushButton(u'导入材质信息')
        self.map_choose_btn.clicked.connect(self._choose)
        self.map_create_btn = QPushButton(u'导入材质ma')
        self.map_create_btn.clicked.connect(self._import_material_ma)
        self.mat_map_layout.addWidget(self.map_label)
        self.mat_map_layout.addWidget(self.map_edit)
        self.mat_map_layout.addWidget(self.map_choose_btn)
        self.mat_map_layout.addWidget(self.map_create_btn)
        self.mat_map_table = MapTable()
        self.btn_layout = QHBoxLayout()
        self.save_btn = QPushButton(u'保存MatMap')
        self.save_btn.clicked.connect(self._save_matmap)
        self.set_mat_btn = QPushButton(u'赋材质')
        self.set_mat_btn.clicked.connect(self._assigned_mat)
        self.btn_layout.addWidget(self.save_btn)
        self.btn_layout.addWidget(self.set_mat_btn)

        self.main_layout.addWidget(self.main_title)
        self.main_layout.addLayout(self.mat_map_layout)
        self.main_layout.addWidget(self.mat_map_table)
        self.main_layout.addLayout(self.btn_layout)
        self.setLayout(self.main_layout)

    def _choose(self):
        openfile_name = QFileDialog.getOpenFileName(self, u'选择文件', '', 'MatMap(*.matmap;)')
        if openfile_name:
            self.map_edit.setText(openfile_name[0])
            try:
                self._create_map()
            except Exception as e:
                QMessageBox.information(None, u'错误', str(e))

    def _create_map(self):
        map_path = self.map_edit.text()

        if map_path:
            mats = ''
            with open(map_path, "r") as r:
                mats = r.read()
            mat_map = eval(mats) if mats else {}
            self.mat_map_table.set_data(mat_map)

    def _import_material_ma(self):
        openfile_name = QFileDialog.getOpenFileName(self, u'选择材质文件', '', 'ma(*.ma)')
        if openfile_name:
            try:
                cmds.file(openfile_name[0], i=True, usingNamespaces=False)
            except Exception as e:
                QMessageBox.information(None, u'错误', str(e))

    def _save_matmap(self):
        save_info = QFileDialog.getSaveFileName(None, u"选择保存路径", ".", "matmap(*.matmap)")
        if save_info:
            try:
                mat_map = self.mat_map_table.get_map()
                file_path = save_info[0]
                with open(file_path, 'w') as w:
                    w.write(str(mat_map))
                choose = cmds.confirmDialog(title='result', message=u'保存成功', icon='information',
                                            button=[u'确认'], defaultButton=u'取消',
                                            dismissString=u'取消')

            except Exception as _save_matmap_ERR:
                print('_save_matmap_ERR:', str(_save_matmap_ERR))

    def get_similar_name(self, transform_name, mat_map):
        """
        获取mat_map中相似的名称
        :param transform_name: 大纲显示名称
        :param mat_map: 物体-材质 映射
        :return:
        """
        result = ''
        for k, v in mat_map.items():
            if k in transform_name:
                result = v

        return result

    def _remove_namespace(self, name):
        """
        清除namespace
        :param name:
        :return:
        """
        return name.split(':')[-1]

    def _assigned_mat(self):
        mat_map = self.mat_map_table.get_map()
        meshs = cmds.ls(type='mesh', l=1)

        err = []
        for m in meshs:
            ts = cmds.listRelatives(m, p=True)
            t = ts[0] if ts else ''
            # mat_name = mat_map.get(t,'')
            mat_name = self.get_similar_name(t, mat_map)
            mat_name = self._remove_namespace(mat_name)
            if mat_name:
                try:
                    sg = mat_name + 'SG'
                    cmds.sets(m, e=1, forceElement=sg)
                except Exception as _assigned_mat_ERR:
                    err.append(_assigned_mat_ERR)
                    print('_assigned_mat_ERR:', str(_assigned_mat_ERR))
        if err:
            QMessageBox.information(None, u'提醒', str(err))
        else:
            QMessageBox.information(None, u'OK', u'完成')




def _run():
    try:
        dialog.show()
        dialog.setHidden(False)
    except :

        dialog = SetMat(parent=maya_main_window())
        dialog.show()
def _doc():
    doc = {'name':u'导入Houdini材质信息','doc':u'导入Houdini材质信息与纯材质ma文件快捷赋材质'}
    return doc

if __name__ == '__main__':
    app =QApplication(sys.argv)
    d = SetMat()
    d.show()
    sys.exit(app.exec_())

